import torch
import torchvision
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets
from torchvision.datasets import ImageFolder
from torchvision.transforms import transforms
from tqdm import tqdm
from common.utils import to_abs_path, ensure_pdir_exist, check_pdir_exist, get_pdir, check_file_exist

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'device: {device}')


def load_resnet34(c):
    if check_file_exist(c.pretrain_model_weight_path):
        model = torchvision.models.resnet34()
        model.load_state_dict(torch.load(c.pretrain_model_weight_path))
        return model
    return torchvision.models.resnet34(weights=torchvision.models.ResNet34_Weights)


class FinetuneResNet34(torch.nn.Module):
    def __init__(self, c):
        super(FinetuneResNet34, self).__init__()
        self.model = load_resnet34(c)
        for i in self.model.parameters():
            i.requires_grad = False
        self.model.fc = torch.nn.Sequential(
            torch.nn.Dropout(c.dropout),
            torch.nn.Linear(self.model.fc.in_features, c.num_classes)
        )

    def forward(self, x):
        return self.model(x)


def test(model, test_loader):
    old_training = model.training
    model.eval()
    total = 0
    correct = 0
    with torch.no_grad():
        for x, y in tqdm(test_loader):
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            out: torch.Tensor = model(x)
            total += x.shape[0]
            correct += (out.argmax(1) == y).sum().item()
    model.train(old_training)
    return correct / total


def draw(c, avg_loss, train_acc, test_acc):
    c.loss_list.append(avg_loss)
    c.train_acc_list.append(train_acc)
    c.test_acc_list.append(test_acc)
    plt.plot(c.train_acc_list, label='train_acc')
    plt.plot(c.test_acc_list, label='test_acc')
    plt.legend()
    ensure_pdir_exist(c.acc_img_path)
    plt.savefig(c.acc_img_path)
    plt.close()
    plt.plot(c.loss_list, label='loss')
    plt.legend()
    ensure_pdir_exist(c.loss_img_path)
    plt.savefig(c.loss_img_path)
    plt.close()


def train(c, train_loader, test_loader):
    model = FinetuneResNet34(c)
    model.to(device)
    model.train()

    optimizer = torch.optim.AdamW(model.parameters(), lr=c.lr, weight_decay=c.weight_decay)
    criterion = torch.nn.CrossEntropyLoss()
    max_test_acc = 0
    for epoch in range(c.epochs):
        total_item = 0
        total_loss = 0.0
        total_correct = 0
        for x, y in tqdm(train_loader):
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            out: torch.Tensor = model(x)
            loss = criterion(out, y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_item += x.shape[0]
            total_loss += loss.item() * x.shape[0]
            total_correct += (out.argmax(1) == y).sum().item()
        avg_loss = total_loss / total_item
        train_acc = total_correct / total_item
        test_acc = test(model, test_loader)
        print(f'epoch: {epoch + 1}/{c.epochs}, loss: {avg_loss:.4f}, '
              f'train_acc: {train_acc:.4f} test_acc: {test_acc:.4f}')
        draw(c, avg_loss, train_acc, test_acc)
        if test_acc > max_test_acc:
            max_test_acc = test_acc
            ensure_pdir_exist(c.model_path)
            torch.save(model.state_dict(), c.model_path)
            print(f'更新max_test_acc={max_test_acc}, save model to {c.model_path}')
            c.no_improve_count = 0
        else:
            c.no_improve_count += 1
            if c.no_improve_count >= c.early_stop:
                print(f'验证集f1没有提升{c.early_stop}次')
                if c.enable_early_stop:
                    print('结束训练')
                    break
    pass


def do():
    c = Config()
    train_loader, test_loader = get_dataloader(c)
    train(c, train_loader, test_loader)
    pass


class Config:
    def __init__(self):
        self.train_path = to_abs_path('./data/Stanford_Cars/train')
        self.test_path = to_abs_path('./data/Stanford_Cars/test')
        # batch_size 的设置考虑训练数据量
        # 学习率考虑 batch_size，同步缩放
        self.train_batch_size = 16
        self.test_batch_size = 128
        self.num_classes = 196
        self.dropout = 0.5
        self.weight_decay = 1e-4
        self.resize = 16 * 32
        self.lr = 3e-4
        self.epochs = 300

        self.model_path = to_abs_path('./models/resnet34_c_car.pth')
        self.pretrain_model_weight_path = to_abs_path('./cache/resnet34_weight.pth')
        self.enable_early_stop = False
        self.early_stop = 6
        self.no_improve_count = 0
        self.loss_list = []
        self.loss_img_path = to_abs_path('./img/resnet34_c_car_loss3.png')
        self.train_acc_list = []
        self.test_acc_list = []
        self.acc_img_path = to_abs_path('./img/resnet34_c_car_acc3.png')


def get_dataloader(c):
    resize = c.resize
    train_transform = transforms.Compose([
        transforms.Resize((resize, resize)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomRotation(25),
        transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.3),
        transforms.RandomGrayscale(p=0.1),
        transforms.ToTensor(),
        transforms.RandomErasing(p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value='random')
    ])
    test_transform = transforms.Compose([
        transforms.Resize((resize, resize)),
        transforms.ToTensor(),
    ])
    train_dataset = ImageFolder(c.train_path, transform=train_transform)
    test_dataset = ImageFolder(c.test_path, transform=test_transform)
    train_loader = DataLoader(train_dataset, batch_size=c.train_batch_size, shuffle=True, num_workers=4)
    test_loader = DataLoader(test_dataset, batch_size=c.test_batch_size, shuffle=False, num_workers=4)
    return train_loader, test_loader


if __name__ == '__main__':
    do()
